import numpy as np
import pandas as pd
from sklearn import metrics
import lightgbm as lgb
from sklearn.model_selection import StratifiedKFold
from sklearn import feature_selection
from sklearn.ensemble import ExtraTreesClassifier

trn_data = pd.read_csv('./data/feas_lgb_merge_300.csv', header=0)
indices = list(range(len(trn_data)))
np.random.shuffle(indices)
trn_data1 = trn_data.iloc[indices, :]

n_train = int(6/7 * len(trn_data1))
X_, Xv_ = trn_data1.iloc[:n_train, :-2], trn_data1.iloc[n_train:, :-2]
y, yv = trn_data1['type'][:n_train], trn_data1['type'][n_train:]
X = X_
Xv = Xv_

clf = ExtraTreesClassifier(n_estimators=50).fit(X_, y)
fs = feature_selection.SelectFromModel(clf, prefit=True)
X = pd.DataFrame(fs.transform(X_))
Xv = pd.DataFrame(fs.transform(Xv_))
print('ori:', X_.shape, 'sel:', X.shape)

features = X.columns
fold = StratifiedKFold(n_splits=6, shuffle=True, random_state=100)
params = {
    'learning_rate': 0.1,
    'n_estimators': 5000,
    'early_stopping_rounds': 100,
    'boosting_type': 'gbdt',
    'objective': 'multiclass',
    'num_class': 3,
    # 'num_leaves': 5,
    'verbosity': -1,
    # 'verbose': -1,
    # 'max_depth': 3,
    # 'feature_fraction': 0.8,    # 特征采样
    # 'class_weight': {0: 0.233, 1: 0.636, 2: 0.131}
}

models = []
oof_v = []
oof_r = []
for index, (train_idx, val_idx) in enumerate(fold.split(X, y)):
    train_set = lgb.Dataset(X.iloc[train_idx], y.iloc[train_idx])
    val_set = lgb.Dataset(X.iloc[val_idx], y.iloc[val_idx])

    model = lgb.train(params, train_set, valid_sets=[train_set, val_set], verbose_eval=100)
    # model = lgb.train(params, train_set)
    models.append(model)

    val_pred = np.argmax(model.predict(X.iloc[val_idx]), axis=1)
    f1_val = metrics.f1_score(y.iloc[val_idx], val_pred, average='macro')
    oof_v.append(f1_val)

    pred_v = model.predict(Xv)
    pred_vc = np.argmax(pred_v, axis=1)
    f1_res = metrics.f1_score(yv, pred_vc, average='macro')
    oof_r.append(f1_res)
    print(index, 'val f1:', f1_val, 'res f1:', f1_res)


f1_val_all, f1_res_all = np.mean(oof_v), np.mean(oof_r)
print('oof val f1:', f1_val_all, 'oof res f1:', f1_res_all)

